from config import  get_config_from_xml
import os
from train import training_main
from data_proc.prepare_dataset import prepare_dataset
from evaluate import evaluate
#五个损失函数
from models.R2_UNet import R2_UNet
from models.UNet import UNet
from models.Res_UNet import ResUNet
from models.Attn_UNet import Attn_UNet

from loss_fn.Jaccard_loss import JaccardLoss
from loss_fn.Tversky_loss import TverskyLoss
from loss_fn.Combined_Loss import CombinedLoss
from segmentation_models_pytorch.losses import DiceLoss
from torch.nn import BCEWithLogitsLoss

# 四个模型
def preSet_main(config_file):
    # 检查配置文件是否存在
    if not os.path.exists(config_file):
        print(f"错误：配置文件 '{config_file}' 不存在。")
        return  # 停止执行后续代码

    cfg = get_config_from_xml(config_file)
    #准备数据划分
    prepare_dataset(cfg)
    model_name = cfg['MODEL_NAME']
    loss_name =  cfg['LOSS_FXN']

    # 实例化模型,将纯字符,转换成模型函数的方法，即文字映射到函数上
    model = globals()[model_name](cfg['IN_CHANNELS'],cfg['OUT_CHANNELS']).to(
        cfg['DEVICE'])
    # 实例化损失函数,将纯字符,转换成损失函数的方法，即文字映射到函数上

    if loss_name == 'DiceLoss' or loss_name == 'CombinedLoss':
        loss_fn = globals()[loss_name](mode='binary')
    else:
        loss_fn = globals()[loss_name]()

    #训练函数
    training_main(cfg, model, loss_fn)
    #预测函数
    evaluate(cfg, batch_size=1)



if "__main__" == __name__:
    config_file = 'configurations/config_attn_unet_dice_epoch50.xml'
    # # preSet_main(config_file)
    # xml_list = ['config_unet_bce_epoch50.xml'
    #             'config_attn_unet_bce_epoch50.xml',\
    #             'config_attn_unet_combined_epoch50.xml',\
    #             'config_attn_unet_dice_epoch50.xml',\
    #             'config_attn_unet_jc_epoch50.xml',\
    #             'config_attn_unet_tversky_epoch50.xml',\
    #             'config_r2_unet_bce_epoch50.xml',\
    #             'config_res_unet_bce_epoch50.xml' ]
    # # config_file = 'configurations/config_attn_unet_dice_epoch40.xml'
        # for xml in xml_list:
        #     config_file=f'configurations/{xml}'
    preSet_main(config_file)


